import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
import sys
import scipy
import statsmodels.api as sm


class MultiFactorModelCovariance():
    
    def regress(self, X,R, N_stocks=50, N_features=5, N_T=300 ):
        '''
        mod: model
        res: fitting result
        spec_ret: specific return on each stock without adjustments
        factor_ret: factor return on panel data of time T.
        X: factor Weight  (N_T,N_features,N_stocks)
        R: return of stock (N_T,N_stocks)
        '''
        

        mod={}
        res={}
        factor_return={}
        spec_ret={}
        # Regress stock return to feature on panel
        for time_step in range(N_T):
            mod[time_step] = sm.OLS(R[time_step,:],X[time_step,:,:])
            res[time_step] = mod[time_step].fit()
            factor_return[time_step] = res[time_step].params
            # ?sm.regression.linear_model.RegressionResults
            spec_ret[time_step]=R[time_step,:]-res[time_step].predict(X[time_step,:,:])


        '''
        tau_sigma_spec_return obtained by Newey-West:
        '''

        tau_specific_ret=100
        spec_ret_nw=pd.DataFrame(spec_ret).T.values


        '''
        tau_corr & tau_sigma obtained by Newey-West test.

        '''

        tau_corr=100
        tau_sigma=10

        tau_specific_ret=100
        df_factor_return=pd.DataFrame(factor_return).T

        '''
        Seperate covariance matrix into correlation matrix and volatility of differnent time dynamics(aka the lags/ half time of autocorrelation are different.)

        then we treat volatility and correlation matrix seperately.
        '''
        corr_mv_window=df_factor_return.ewm(halflife=tau_corr)
        sigma_mv_window=df_factor_return.ewm(halflife=tau_sigma)
        corr_raw=corr_mv_window.corr()
        sigma_raw=np.sqrt(sigma_mv_window.var())

        '''
        estimate specific return (residual return: R=XF+U)
        '''

        df_spec_ret=pd.DataFrame(spec_ret)

        corr_raw_np=corr_raw.values.reshape([N_T,N_features,N_features])
        sigma_raw_np=sigma_raw.values
        ans=np.dot(corr_raw_np,sigma_raw_np.reshape([N_T,N_features,1]))
        ans=corr_raw_np*sigma_raw_np.reshape([N_T,N_features,1])*sigma_raw_np.reshape([N_T,1,N_features])
        robust_cov=ans
        #cov=corr_mv_window.cov().values.reshape([N_T,N_features,N_features])

        '''
        compute XFX^T for each time period.

        we apply einstein summation on X_{tnk} and F_{tkk}
        '''
        tmp1=np.einsum('tnk,tkk->tnk', X, robust_cov)
        '''
        Here l=n.
        '''
        xfx=np.einsum('tnk,tkl->tnl', X, np.transpose(tmp1,(0,2,1)))

        ui=np.zeros([N_T,N_stocks,N_stocks])
        for timestamp in range(N_T):
            ui[timestamp,:,:]=np.diag(spec_ret_nw[timestamp,:])

        model_covariance=xfx+ui
        # np.save("./data/covariance_matrix.npy", model_covariance)
        # np.save("./data/return.npy",R)
        return model_covariance[-1,:,:]